#include "StdAfx.h"
#include "Resources.h"

enum class ExecutionMode_t {
	FsScan,
	PsScan
};

int32_t wmain(int32_t nArgc, const wchar_t* pArgv[]) {
	vector<wstring> Args(&pArgv[0], &pArgv[0 + nArgc]);
	uint64_t qwFilterLackingFlags = 0;
	uint64_t qwFilterOnlyFlags = 0;
	wstring TargetPath;
	uint32_t dwPid = 0;
	ExecutionMode_t Mode;
	shared_ptr<Interface> Intf = make_shared<Interface>(Args);
	shared_ptr<ApiTable> DynamicApis = make_shared<ApiTable>(Intf);

	if (nArgc >= 3) {
		for (vector<wstring>::const_iterator i = Args.begin(); i != Args.end(); ++i) {
			wstring Arg = *i;
			transform(Arg.begin(), Arg.end(), Arg.begin(), ::tolower);

			if (Arg == L"-p") {
				dwPid = _wtoi((*(i + 1)).c_str());
			}
			else if (Arg == L"--mode") {
				if (*(i + 1) == L"fs") {
					Mode = ExecutionMode_t::FsScan;
				}
				else if (*(i + 1) == L"ps") {
					Mode = ExecutionMode_t::PsScan;
				}
			}
			else if (Arg == L"--filter-only-lacking") {
				for (vector<wstring>::const_iterator SearchItr = i; SearchItr != Args.end(); ++SearchItr) {
					wstring MitigationArg = *SearchItr;
					transform(MitigationArg.begin(), MitigationArg.end(), MitigationArg.begin(), ::tolower);

					if (MitigationArg == L"aslr") {
						qwFilterLackingFlags |= SEARCH_FLAG_ASLR;
					}
					else if (MitigationArg == L"seh") {
						qwFilterLackingFlags |= SEARCH_FLAG_SECURE_EXCEPTIONS;
					}
					else if (MitigationArg == L"cfg") {
						qwFilterLackingFlags |= SEARCH_FLAG_CFG;
					}
					else if (MitigationArg == L"dep") {
						qwFilterLackingFlags |= SEARCH_FLAG_DEP;
					}
					else if (MitigationArg == L"acg") {
						qwFilterLackingFlags |= SEARCH_FLAG_ACG;
					}
					else if (MitigationArg == L"cig") {
						qwFilterLackingFlags |= SEARCH_FLAG_CIG;
					}
					else if (MitigationArg == L"*") {
						qwFilterLackingFlags = -1;
					}
				}
			}
			else if (Arg == L"--filter-for-active") {
				for (vector<wstring>::const_iterator SearchItr = i; SearchItr != Args.end(); ++SearchItr) {
					wstring MitigationArg = *SearchItr;
					transform(MitigationArg.begin(), MitigationArg.end(), MitigationArg.begin(), ::tolower);

					if (MitigationArg == L"aslr") {
						qwFilterOnlyFlags |= SEARCH_FLAG_ASLR;
					}
					else if (MitigationArg == L"seh") {
						qwFilterOnlyFlags |= SEARCH_FLAG_SECURE_EXCEPTIONS;
					}
					else if (MitigationArg == L"cfg") {
						qwFilterOnlyFlags |= SEARCH_FLAG_CFG;
					}
					else if (MitigationArg == L"dep") {
						qwFilterOnlyFlags |= SEARCH_FLAG_DEP;
					}
					else if (MitigationArg == L"acg") {
						qwFilterOnlyFlags |= SEARCH_FLAG_ACG;
					}
					else if (MitigationArg == L"cig") {
						qwFilterOnlyFlags |= SEARCH_FLAG_CIG;
					}
					else if (MitigationArg == L"*") {
						qwFilterOnlyFlags = -1;
					}
				}
			}
			else if (Arg == L"--path") {
				TargetPath = (*(i + 1));
			}
		}

		switch (Mode) {
		case ExecutionMode_t::PsScan: {
			if (GrantSelfSePrivilege(L"SeDebugPrivilege")) {
				Intf->Log("... successfully granted SeDebug privilege to self\r\n");
			}
			else {
				Intf->Log("... failed to grant SeDebug privilege to self\r\n");
			}
			
			EnumPsMitigations(Intf, DynamicApis, qwFilterLackingFlags, qwFilterOnlyFlags);
			break;
		}
		case ExecutionMode_t::FsScan: {
			if (!TargetPath.empty() && qwFilterLackingFlags) {
#ifndef _WIN64
				HMODULE hKernel32Module = GetModuleHandleW(L"Kernel32.dll");
				Wow64DisableWow64FsRedirection_t Wow64DisableWow64FsRedirection = (Wow64DisableWow64FsRedirection_t)GetProcAddress(hKernel32Module, "Wow64DisableWow64FsRedirection");
				void* pOldValue = nullptr;

				if (Wow64DisableWow64FsRedirection != nullptr) {
					if (Wow64DisableWow64FsRedirection(&pOldValue)) {
						Intf->Log("... successfully disabled Wow64 FS redirection\r\n");
					}
					else {
						Intf->Log("... failed to disable Wow64 FS redirection\r\n");
					}
				}
#endif
				Intf->Log("... starting scan of %ws...\r\n", TargetPath.c_str());

				uint64_t qwStartTick = GetTickCount64();

				if (GetFileAttributesW(TargetPath.c_str()) & FILE_ATTRIBUTE_DIRECTORY) {
					uint32_t dwTotalPe = 0;
					vector<MitigationlessPe> PeVec = MitigationlessPe::ScanFiles(Intf, DynamicApis, TargetPath, qwFilterLackingFlags, dwTotalPe);

					for (vector<MitigationlessPe>::const_iterator Itr = PeVec.begin(); Itr != PeVec.end(); ++Itr) {
						Itr->Print();
					}

					Intf->Log("... %d total non-mitigation PE files out of %d total PE (%f%%)\r\n", PeVec.size(), dwTotalPe, (float)PeVec.size() / dwTotalPe * 100.0);
				}
				else {
					try {
						MitigationlessPe Pe(Intf, DynamicApis, TargetPath, qwFilterLackingFlags);
						Pe.Print();
					}
					catch (int32_t nError) {
						if (nError == 3) {
							Intf->Log("... PE at %ws has none of the specified weak exploit mitigation features\r\n", TargetPath.c_str());
						}
						else {
							Intf->Log("... invalid PE file at %ws (does not exist or does not contain code)\r\n", TargetPath.c_str());
						}
					}
				}

				float fElapsedTime = GetTickCount64() - qwStartTick;
				Intf->Log("... scan completed (%f second duration)\r\n", fElapsedTime / 1000.0);
			}
			else {
				Intf->Log("... fatal error: a target path and minimum of one mitigation filter criteria are required in this execution mode.\r\n");
			}

			break;
		}
		default: break;
		}
	}
	else {
		HMODULE	hSelfModule = GetModuleHandleA(nullptr);
		HRSRC hResourceInfo;
		HGLOBAL hResourceData;
		char* pRsrcData = nullptr;
		uint32_t dwRsrcSize;

		if ((hResourceInfo = FindResourceA(hSelfModule, IDR_USAGE_TEXT_NAME, RT_RCDATA))) {
			if ((hResourceData = LoadResource(hSelfModule, hResourceInfo))) {
				dwRsrcSize = SizeofResource(hSelfModule, hResourceInfo);
				pRsrcData = (char*)LockResource(hResourceData);
				unique_ptr<uint8_t[]> RsrcBuf = make_unique<uint8_t[]>(dwRsrcSize + 1); // Otherwise the resource text may bleed in to the rest of the .rsrc section
				memcpy(RsrcBuf.get(), pRsrcData, dwRsrcSize);
				Intf->Log("%s\r\n", RsrcBuf.get());
			}
		}
	}

	return 0;
}